/*msg_key_large = SHA256 (substr (key, 88+x, 32) + plaintext + random_padding);
msg_key = substr (msg_key_large, 8, 16);
sha256_a = SHA256 (msg_key + substr (key, x, 36));
sha256_b = SHA256 (substr (key, 40+x, 36) + msg_key);
aes_key = substr (sha256_a, 0, 8) + substr (sha256_b, 8, 16) + substr (sha256_a, 24, 8);
aes_iv = substr (sha256_b, 0, 8) + substr (sha256_a, 8, 16) + substr (sha256_b, 24, 8);
 */
use base64;
use std::convert::TryInto;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpStream, tcp::ReadHalf, tcp::WriteHalf};
use rand_core::OsRng;
use x25519_dalek::{EphemeralSecret, PublicKey};
use sha2::{Sha256, Digest};
use rand::Rng;
use openssl::aes::{AesKey, aes_ige};
use openssl::symm::Mode;

fn sha256(items : &[&[u8]]) -> Vec<u8> {
    
    let mut hasher = Sha256::new();
    for item in items {
	hasher.update(item);
    }
    let result = hasher.finalize();

    return result[..].to_vec();
}

fn encrypt_message(key : &[u8], plain_text : &str) -> String {

    let mut payload : Vec<u8> = Vec::new();
    let bs = u32::to_be_bytes(plain_text.as_bytes().len() as u32);
    payload.extend(bs);
    payload.extend(plain_text.as_bytes());
    let padding_size = 64 - (payload.len() % 16);
    let padding = (0..padding_size).map(|x|{let b : u8 = rand::thread_rng().gen(); b}).collect::<Vec<_>>();
    let padding_partial = &padding[padding.len()-32..];
    payload.extend(&padding[..]);

    let sha256_a = sha256(&[key, padding_partial]);
    let sha256_b = sha256(&[padding_partial, key]);

    let mut b_aes_key : Vec<u8> = Vec::new();
    b_aes_key.extend(&sha256_a[0..8]);
    b_aes_key.extend(&sha256_b[8..24]);
    b_aes_key.extend(&sha256_a[24..32]);

    let mut b_aes_ige : Vec<u8> = Vec::new();
    b_aes_ige.extend(&sha256_b[0..8]);
    b_aes_ige.extend(&sha256_a[8..24]);
    b_aes_ige.extend(&sha256_b[24..32]);

    //println!("b_aes_key: length = {}, data = {:?}", b_aes_key.len(), &b_aes_key);
    let aes_key = AesKey::new_encrypt(&b_aes_key).unwrap();
    let mut encrypted = vec![0;payload.len()];
    aes_ige(&payload, &mut encrypted, &aes_key, &mut b_aes_ige, Mode::Encrypt);

    let mut message : Vec<u8> = Vec::new();
    message.extend(padding_partial);
    message.extend(encrypted);

    String::from(base64::encode(message) + "\n")
}

fn decrypt_message(key : &[u8], encrypted_message : &str) -> String {
    let encrypted_message = base64::decode(encrypted_message.trim()).unwrap();
    let padding_partial = &encrypted_message[..32];
    let sha256_a = sha256(&[key, padding_partial]);
    let sha256_b = sha256(&[padding_partial, key]);

    let encrypted = &encrypted_message[32..];

    let mut b_aes_key : Vec<u8> = Vec::new();
    b_aes_key.extend(&sha256_a[0..8]);
    b_aes_key.extend(&sha256_b[8..24]);
    b_aes_key.extend(&sha256_a[24..32]);

    let mut b_aes_ige : Vec<u8> = Vec::new();
    b_aes_ige.extend(&sha256_b[0..8]);
    b_aes_ige.extend(&sha256_a[8..24]);
    b_aes_ige.extend(&sha256_b[24..32]);

    let aes_key = AesKey::new_decrypt(&b_aes_key).unwrap();
    let mut message : Vec<u8> = Vec::new();
    let mut decrypted = vec![0; encrypted.len()];
    aes_ige(encrypted, &mut decrypted, &aes_key, &mut b_aes_ige, Mode::Decrypt);

    let length : u32 = u32::from_be_bytes(decrypted[..4].try_into().unwrap());
    let bs = (&decrypted[4..4+length as usize]).to_vec();

    String::from_utf8(bs).unwrap()
}

async fn exchange_key(reader : &mut BufReader<ReadHalf<'_>>, writer : &mut WriteHalf<'_>) -> Vec<u8> {
    let mut line = String::new();
    let local_secret = EphemeralSecret::new(OsRng);
    let local_public = PublicKey::from(&local_secret);
    let encoded = base64::encode(local_public.as_bytes()) + "\n";
    writer.write_all(encoded.as_bytes()).await.unwrap();
    reader.read_line(&mut line).await.unwrap();
    let decoded = base64::decode(line.trim()).unwrap();
    line.clear();
    let remote_public : [u8; 32] = decoded.try_into().unwrap();
    let remote_public = PublicKey::from(remote_public);
    let local_shared_secret = local_secret.diffie_hellman(&remote_public);
    println!(
        "len = {}, shared_secret = {:?}",
        local_shared_secret.as_bytes().len(),
        local_shared_secret.as_bytes()
    );

    return local_shared_secret.as_bytes().to_vec();
}

#[tokio::main]
async fn main() {

    let mut console_reader = BufReader::new(tokio::io::stdin());
    let mut console_buf = String::new();

    let mut conn = TcpStream::connect("127.0.0.1:8080").await.unwrap();
    let (mut reader, mut writer) = conn.split();

    let mut reader = BufReader::new(reader);
    let mut line = String::new();

    let key = exchange_key(&mut reader, &mut writer).await;

    loop {
        tokio::select! {
            result = console_reader.read_line(&mut console_buf) => {
		let encoded = encrypt_message(&key, &console_buf);
		writer.write_all(encoded.as_bytes()).await.unwrap();
		console_buf.clear();
            },
            result = reader.read_line(&mut line) => {
		let bytes_read = result.unwrap();
                if bytes_read == 0 {
                    break;
                }
		let decoded = decrypt_message(&key, &line.trim());
		println!("Received: {}", decoded.trim());
		line.clear();
            },
        };
    }


}
